"""
https://www.tensorflow.org/guide/keras/train_and_evaluate
"""

import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow import keras
from tensorflow.keras import layers, losses, metrics, optimizers, Model
from python_ai.common.xcommon import *
import os
import sys
from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt

# data
x, y = load_boston(return_X_y=True)
scaler = StandardScaler()
x = scaler.fit_transform(x)
y = scaler.fit_transform(y.reshape([-1, 1]))

x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=777)
m, n = x_train.shape
sep((m, n))

# Reserve 10,000 samples for validation
n_val = int(np.ceil(m * 0.1))
x_val = x_train[-n_val:]
y_val = y_train[-n_val:]
x_train = x_train[:-n_val]
y_train = y_train[:-n_val]

# model
alpha = 0.001
L1 = 400
L2 = 200
n_epoch = 10
batch_size = 128
inputs = layers.Input(shape=(n,), dtype=tf.float32, name='Inputs')
dense1 = layers.Dense(L1, tf.nn.relu, dtype=tf.float32, name='Dense1')(inputs)
dense2 = layers.Dense(L2, tf.nn.relu, dtype=tf.float32, name='Dense2')(dense1)
outputs = layers.Dense(1, None, dtype=tf.float32, name='Outputs')(dense2)
model = Model(inputs=inputs, outputs=outputs)

model.compile(
    optimizer=optimizers.RMSprop(learning_rate=alpha),
    loss=losses.mean_squared_error,
    metrics=[tfa.metrics.r_square.RSquare(y_shape=(1, ))]
)

history = model.fit(
    x_train,
    y_train,
    batch_size=batch_size,
    epochs=n_epoch,
    validation_data=(x_val, y_val)
)
print(history.history)

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(history.history['loss'], label='train')
plt.plot(history.history['val_loss'], label='validation')
plt.legend()
plt.title('loss')
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(history.history['r_square'], label='train')
plt.plot(history.history['val_r_square'], label='validation')
plt.legend()
plt.title('r2')


r = model.evaluate(x_test, y_test)
print(r)

